#include <iostream>
#include <stdlib.h>
#include <time.h>
#include "qcircuit.h"

qcircuit::qcircuit(qreg & qr, creg & cr):nqr(qr.n), ncr(cr.n){
    n = nqr + ncr;
    p_qr = & qr;
    p_cr = & cr;
    supo = new qubit;
    // p_qstate = new string[nqr + ncr];
    p_qstate = new string[nqr + 1];
    for ( int i = 0; i <= nqr; ++ i )
        p_qstate[i] =  "";

    // 叠加态，对 nqr 个量子比特做张量积
    *supo = p_qr->reg[0];
    for ( int i = 1; i < nqr; ++ i ){
        *supo = tensorProduct(*supo, p_qr->reg[i]);
    }

    for ( int i = 0; i < nqr; ++ i ){
        p_qstate[i] = "q";
        p_qstate[i] += char('0' + i);
        p_qstate[i] += " ---";
    }

    p_qstate[nqr] = "c";
    p_qstate[nqr] += char('0'+ncr-1);
    p_qstate[nqr] += " ===";
}

bool dicing(double p){
    double n = 1.0 * rand() / RAND_MAX;
    if ( n > p ){
        return false;
    }else{
        return true;
    }
}

string _to_string(int num, unsigned wide){
    string ket = "";
    char s[10];
    itoa(num, s, 2);
    for ( unsigned i = 0; i < wide - strlen(s); ++ i )
        ket += '0';
    for ( unsigned i = 0; i < strlen(s); ++ i )
        ket += s[i];
    return ket;
}

qcircuit::~qcircuit(){
    srand((unsigned)time(0));
    for ( int i = 0; i < ncr; ++ i ){
        if ( p_cr->reg[i] != 3 ){
            cout << "q" << i << ": " << 1 - p_cr->reg[i] << ' '<< p_cr->reg[i] << endl;
        }
    }
    int cnt[(1<<ncr)] = {0};
    int idx = 0;
    for ( int i = 0; i < MAX; ++ i ){
        idx = 0;
        for ( int j = 0; j < ncr; ++ j ){
            // p_cr->reg[j]
            if ( p_cr->reg[j] != 3 ){
                if ( dicing(p_cr->reg[j]) ){
                    idx = (idx << 1) + 1;
                }else{
                    idx = (idx << 1) + 0;
                }
            }
        }
        ++ cnt[idx];
    }
    printf("{\n");
    for ( int i = 0; i < (1<<ncr); ++ i ){
        if ( cnt[i] )
            printf("\t'%s': %d\n", _to_string(i, ncr).c_str(), cnt[i]);
    }
    printf("}\n");
    delete [] p_qstate;
    delete [] supo;
}

void qcircuit::h(qubit & q){
    int idx = &q - &p_qr->reg[0];
    
    gate Hgate(2,2,H);
    gate Igate(2,2,I);
    gate g; // default (2,2,I);
    if ( idx == 0 )
        g = Hgate;  // I kron ... H kron I ...
    for ( int i = 1; i < nqr; ++ i ){
        if ( i == idx )
            g = tensorProduct(g, Hgate);
        else
            g = tensorProduct(g, Igate);
    }
    linearTransform(g, *supo);
    // 更新电路视图
    if ( p_qstate[idx][p_qstate[idx].size()-4] == '-' ){
        p_qstate[idx][p_qstate[idx].size()-4] = 'H';
        return;
    }
    for ( int i = 0; i < nqr; ++ i )
        p_qstate[i] += ( ( & p_qr->reg[i] == & q ) ? "H---" : "----" );
    p_qstate[nqr] += "====";
    return;
}

void qcircuit::cx(qubit & q1, qubit & q2){
    int idx1 = &q1 - &p_qr->reg[0];
    int idx2 = &q2 - &p_qr->reg[0];
    // 交换系数
    complex<type> tmp;
    for ( int i = 0; i < supo->n; ++ i ){
        if ( (( i >> (nqr - idx1 - 1) ) & 1) == 1 ){
            if ( (( i >> (nqr - idx2 - 1) ) & 1) == 1 ){
                int j = i - pow(2, nqr - idx2 - 1);
                tmp = supo->v[j];
                supo->v[j] = supo->v[i];
                supo->v[i] = tmp;
            }
        }
    }
    // 更新电路视图
    for ( int i = 0; i < nqr; ++ i ){
        if ( & p_qr->reg[i] == & q1 ){
            p_qstate[i] += "*---";
        }else if ( & p_qr->reg[i] == & q2 ){
            p_qstate[i] += "x---";
        }else if ( min(&q1, &q2) < & p_qr->reg[i] && & p_qr->reg[i] < max(&q1, &q2) ){
            p_qstate[i] += "|---";
        }else{
            p_qstate[i] += " ---";
        }
    }
    p_qstate[nqr] += "====";
}

void qcircuit::measure(qubit & q, double & c){
    int idx1 = &q - &p_qr->reg[0];
    int idx2 = &c - &p_cr->reg[0];
    
    // 计算 q = |1> 的概率
    c = 0;
    for ( int i = 0; i < supo->n; ++ i ){
        if ( (( i >> (nqr - idx1 - 1) ) & 1) == 1 ){
            c += supo->v[i].real() * supo->v[i].real() + supo->v[i].imag() * supo->v[i].imag();
        }
    }

    //supo->getDim();
    //supo->getAmplitude();

    // 更新电路视图
    for ( int j = 0; j < nqr; ++ j )
        if ( j > idx1 )
            p_qstate[j] += "|---";
        else if ( j == idx1 )
            p_qstate[j] += "M---";
        else
            p_qstate[j] += "----";
    p_qstate[nqr] += "v";
    p_qstate[nqr] += char('0'+idx2);
    p_qstate[nqr] += "==";
}

void qcircuit::draw(){
    cout << "Draw: " << endl;
    for ( int i = 0; i <= nqr; ++ i )
        cout << p_qstate[i] << endl;
    cout << "Draw done! " << endl;
}

// void qcircuit::getAmplitude(){
//     for ( int i = 0; i < nqr; ++ i ){
//         cout << 'q' << i << ": ";
//         p_qr->reg[i].getAmplitude();
//     }
// }

